{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import pickle\n",
    "import time\n",
    "import sklearn.metrics.pairwise as smp\n",
    "import random\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "UnpicklingError",
     "evalue": "invalid load key, '\\xce'.",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mUnpicklingError\u001b[0m                           Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-4-1dc35dcfdc59>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m()\u001b[0m\n\u001b[0;32m      2\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      3\u001b[0m \u001b[1;32mwith\u001b[0m \u001b[0mopen\u001b[0m \u001b[1;33m(\u001b[0m\u001b[1;34m'client_0_round_0.pkl'\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m'rb'\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0mf\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 4\u001b[1;33m     \u001b[0mlocal_para\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mpickle\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mload\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mf\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      5\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      6\u001b[0m \u001b[1;32mwith\u001b[0m \u001b[0mopen\u001b[0m \u001b[1;33m(\u001b[0m\u001b[1;34m'client_1_round_0.pkl'\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m'rb'\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0mf\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mUnpicklingError\u001b[0m: invalid load key, '\\xce'."
     ]
    }
   ],
   "source": [
    "local_para = []\n",
    "\n",
    "with open ('client_0_round_0.pkl', 'rb') as f:\n",
    "    local_para.append(pickle.load(f))\n",
    "\n",
    "with open ('client_1_round_0.pkl', 'rb') as f:\n",
    "    local_para.append(pickle.load(f))\n",
    "    \n",
    "with open ('client_2_round_0.pkl', 'rb') as f:\n",
    "    local_para.append(pickle.load(f))\n",
    "\n",
    "\n",
    "with open ('local.pkl', 'wb') as f:\n",
    "    pickle.dump(local_para, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[8.38292807340622, 8.38292807340622]\n"
     ]
    }
   ],
   "source": [
    "def to_matrix(local_para_list: list) -> dict:\n",
    "    para_matrix = {}\n",
    "    for key in local_para_list[0]:\n",
    "        grad_list = []\n",
    "        for local_para in local_para_list:\n",
    "            grad_list.append(torch.squeeze(local_para[key].reshape(1, -1)))\n",
    "        para_matrix[key] = torch.tensor(np.array([item.cpu().detach().numpy() for item in grad_list]))\n",
    "    return para_matrix\n",
    "\n",
    "def compute_cosine_similarity(local_para_matrix: dict) -> list:\n",
    "    maxcs = []\n",
    "    for key, matrix in local_para_matrix.items():\n",
    "        cs = smp.cosine_similarity(matrix) - np.eye(matrix.shape[0])\n",
    "        maxcs.append(np.max(cs, axis=1))\n",
    "\n",
    "    scores = []\n",
    "    for id in range(len(maxcs[0])):\n",
    "        sum = 0\n",
    "        for key in range(len(maxcs)):\n",
    "            sum += maxcs[key][id]\n",
    "        scores.append(sum)\n",
    "\n",
    "    return scores\n",
    "\n",
    "def select_and_aggregate(scores: list, client_addresses: list, local_para: list, global_para: dict, p_solo: float, p_team: float) -> dict:\n",
    "    score_address = {}\n",
    "    score_para = {}\n",
    "    client_selected = []\n",
    "    para_selected = []\n",
    "\n",
    "    for i, score in enumerate(scores):\n",
    "        score_address[score] = client_addresses[i]\n",
    "        score_para[score] = local_para[i]\n",
    "    score_sorted = score.sort()\n",
    "    score_selected = score_sorted[int(p_solo * len(score)) + 1 : int((1 - p_team) * len(score))]\n",
    "\n",
    "    for score in score_selected:\n",
    "        if score in score_address.values:\n",
    "            client_selected.append(score_address[score])\n",
    "            para_selected.append(score_para[score])\n",
    "\n",
    "    num_selected = len(para_selected)\n",
    "    for idx in range(num_selected):\n",
    "        noise_para = para_selected[idx]\n",
    "        if idx == 0:\n",
    "            for key in noise_para:\n",
    "                global_para[key] = noise_para[key] / num_selected\n",
    "        else:\n",
    "            for key in noise_para:\n",
    "                global_para[key] += noise_para[key] / num_selected\n",
    "    return global_para\n",
    "    \n",
    "    \n",
    "with open('local.pkl', 'rb') as f:\n",
    "    local_para_list = pickle.load(f)\n",
    "\n",
    "local_para_matrix = to_matrix(local_para_list)\n",
    "\n",
    "scores = compute_cosine_similarity(local_para_matrix)\n",
    "\n",
    "# global_para = select_and_aggregate(scores, client_addresses, local_para, global_para, 1/3, 1/3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0 5 2]\n"
     ]
    }
   ],
   "source": [
    "proportion = 1 / 2\n",
    "\n",
    "nets = {'a':0, 'b':1, 'c':2, 'd':3, 'e':4, 'f':5}\n",
    "\n",
    "s = np.random.permutation(len(nets))[:math.floor(proportion*len(nets))]\n",
    "\n",
    "print(s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([1])\n",
      "[9 1 3 0 6 4 8 5 2 7]\n",
      "[9]\n",
      "[5 2 7]\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import random\n",
    "import math\n",
    "\n",
    "target = [5]\n",
    "\n",
    "target = torch.tensor(list(map(lambda x: random.randint(0, 9), target)))\n",
    "\n",
    "print(target)\n",
    "\n",
    "shuffle = np.random.permutation(10)\n",
    "solo_selected = shuffle[:math.floor(0.1*10)]\n",
    "team_selected = shuffle[-math.floor(0.3*10):]\n",
    "\n",
    "print(shuffle)\n",
    "print(solo_selected)\n",
    "print(team_selected)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3\n"
     ]
    }
   ],
   "source": [
    "# solo_selected = np.array([2, 3, 4])\n",
    "\n",
    "print(' '.join(list(map(lambda x: str(x), solo_selected))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0 1 2 3 4 5 6 7 8 9]\n",
      "[2 3 4 5 6]\n"
     ]
    }
   ],
   "source": [
    "p_solo = 0.2\n",
    "p_team = 0.3\n",
    "scores = np.arange(10)\n",
    "print(scores)\n",
    "\n",
    "score_selected = scores[math.ceil(p_solo * len(scores)) : math.floor((1 - p_team) * len(scores))]\n",
    "print(score_selected)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0 1 7 5 9]\n"
     ]
    }
   ],
   "source": [
    "arr = [i for i in range(10)]\n",
    "\n",
    "a = random.sample(list(arr[:3]), 2)\n",
    "\n",
    "b = random.sample(list(arr[3:8]), 2)\n",
    "\n",
    "c = random.sample(list(arr[8:]), 1)\n",
    "\n",
    "b.extend(c)\n",
    "\n",
    "selected = np.append(a, b)  #选中的客户端\n",
    "\n",
    "print(selected)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{0: 99, 1: 92, 2: 91}\n"
     ]
    }
   ],
   "source": [
    "id_score = {0: 99, 1: 92, 2: 91}\n",
    "\n",
    "id_score_sort = sorted(id_score.items(),key = lambda x:x[1],reverse = True)\n",
    "\n",
    "id_score_sort = dict(id_score_sort)\n",
    "print(id_score_sort)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "FedAvg",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
